import torch
import torchvision
from torch import nn
from net import MyAlexNet
import numpy as np
from torch.optim import lr_scheduler
import os

from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

import matplotlib.pyplot as plt

# 解决乱码。中文显示问题
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 导入数据
ROOT_TRAIN = r"D:/PyCharm Coder Files/AlexNet/data/train"
ROOT_TEST = r"D:/PyCharm Coder Files/AlexNet/data/val"

# device : GPU or CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

# 将图像的像素值[0,1]归一化到[-1,1]之间
# 图像预处理之：归一化数据转换。RGB三通道进行 如(R-mean)/std
# 经过这样处理后的数据符合标准正态分布，即均值为0，标准差为1。使模型更容易收敛
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

# 数据预处理
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomVerticalFlip(),  # 随机垂直翻转，来进行数据增强处理
    transforms.ToTensor(),
    normalize])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize])

# 加载数据集
train_dataset = ImageFolder(root=ROOT_TRAIN, transform=train_transform)
val_dataset = ImageFolder(root=ROOT_TEST, transform=val_transform)

train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=32, shuffle=True)

# 加载模型
model = MyAlexNet().to(device)

# 定义一个损失函数
loss_fn = nn.CrossEntropyLoss()

# 定义一个优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 定义学习率:每隔10轮，变为原来的0.05
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)


# 定义一个训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0
    for _, (x, y) in enumerate(dataloader):
        # 前向传播
        image, y = x.to(device), y.to(device)
        output = model(image)
        _, pred = torch.max(output, axis=1)
        cur_loss = loss_fn(output, y)
        cur_acc = torch.sum(y == pred) / output.shape[0]

        # 反向传播
        optimizer.zero_grad()
        cur_loss.backward()
        optimizer.step()

        loss += cur_loss.item()
        current += cur_acc.item()
        n = n + 1

    train_loss = loss / n
    train_acc = current / n
    print("train_loss" + str(train_loss))
    print("train_acc" + str(train_acc))
    return train_loss, train_acc


def val(dataloader, model, loss_fn):
    model.eval()
    loss, current, n = 0.0, 0.0, 0
    with torch.no_grad():
        for _, (x, y) in enumerate(dataloader):
            image, y = x.to(device), y.to(device)
            output = model(image)
            _, pred = torch.max(output, axis=1)
            cur_loss = loss_fn(output, y)
            cur_acc = torch.sum(y == pred) / output.shape[0]

            loss += cur_loss.item()
            current += cur_acc.item()
            n = n + 1

    val_loss = loss / n
    val_acc = current / n
    print("val_loss" + str(val_loss))
    print("val_acc" + str(val_acc))
    return val_loss, val_acc


# 定义画图函数
def matplot_loss(train_loss, val_loss):
    plt.plot(train_loss, label='train_loss')
    plt.plot(val_loss, label='val_loss')
    plt.legend(loc='best')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.title("训练集和验证集loss值对比图")
    plt.show()


def matplot_loss(train_acc, val_acc):
    plt.plot(train_acc, label='train_acc')
    plt.plot(val_acc, label='val_acc')
    plt.legend(loc='best')
    plt.xlabel('epoch')
    plt.ylabel('accuracy')
    plt.title("训练集和验证集accuracy值对比图")
    plt.show()


# 开始训练
loss_train = []
acc_train = []
loss_val = []
acc_val = []

epoch = 20
min_acc = 0
for t in range(epoch):

    print(f"epoch{t + 1}\n---------------------------")
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer)
    val_loss, val_acc = val(val_dataloader, model, loss_fn)

    lr_scheduler.step()

    loss_train.append(train_loss)
    acc_train.append(train_acc)
    loss_val.append(val_loss)
    acc_val.append(val_acc)

    # 只保存最好的模型权重
    if val_acc > min_acc:
        folder = "save_model"
        if not os.path.exists(folder):
            os.mkdir(folder)
        min_acc = val_acc
        print(f"save best model,第{t + 1}轮")
        torch.save(model.state_dict(), 'save_model/best_model.pth')
    if t == epoch - 1:
        torch.save(model.state_dict(), 'save_model/best_model.pth')
    print("\n")

print("Done！")
